import numpy as np
import pickle
import matplotlib.pyplot as plt
import os

def load_and_save_dataset_images(file_name):
    '''
    从 .pkl 文件中加载数据集，并将每个类别中的第一张图像保存为 PNG 文件。
    '''
    try:
        with open(file_name, 'rb') as f:
            # 数据集的顶级数据类型是字典
            raw_dataset = pickle.load(f)

        # 打印数据集的结构信息
        print(f"成功加载文件: {file_name}")
        print(f"数据集的顶级数据类型: {type(raw_dataset)}")
        print(f"数据集包含 {len(raw_dataset)} 个类别。")
        print("---")
        
        # 创建输出文件夹
        output_dir = './output/training_dataset'
        os.makedirs(output_dir, exist_ok=True)
        print(f"输出文件夹已创建: {output_dir}")

        # 遍历字典的键值对，键是类别标签，值是图像列表
        for class_label, images_list in raw_dataset.items():
            print(f"处理类别: {class_label}")

            # 检查images_list的数据类型
            if isinstance(images_list, list) and len(images_list) > 0:
                first_image = images_list[0]
                
                # 创建一个Matplotlib图形对象，但不需要显示
                plt.imshow(first_image, cmap='gray')
                plt.title(f"类别: {class_label}")
                
                # 定义保存路径和文件名
                file_path = os.path.join(output_dir, f"label_{class_label}.png")
                
                # 保存图像
                plt.savefig(file_path)
                print(f"  - 图像已保存至: {file_path}")
                
                # 关闭当前图形，以便下一张图片可以独立创建
                plt.close()
            else:
                print(f"  - 类别 {class_label} 中没有找到图像列表或图像为空。跳过。")
            
            print("---")

    except FileNotFoundError:
        print(f"错误: 文件 {file_name} 未找到。请确保文件路径正确。")
    except Exception as e:
        print(f"加载或保存文件时发生错误: {e}")
        print("请检查数据集结构是否与代码预期一致。")

if __name__ == '__main__':
    # 明确指定要加载的文件
    file_to_load = './training_dataset.pkl'
    load_and_save_dataset_images(file_to_load)